package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{StringIndexEncoder, StringIndexEncoderParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.scalatest.FunSuite

/**
  * @Author tangjinyuan
  * @Date 2019/3/27 14:32
  * @Description
  * @Reviewer
  */
class StringIndexEncoderSuite extends FunSuite {
  val paramStr: String =
    """
      |{
      |	"inputCol": "category",
      |	"outputCol": "categoryIndex",
      |	"handleInvalid": "skip"
      |}
    """.stripMargin

  val sie = new StringIndexEncoder()
  val param: StringIndexEncoderParam = sie.parseParam(new ToolClass().encrypt(paramStr))

  test("sie parameter") {
    assert("category".equals(param.inputCol) && "categoryIndex".equals(param.outputCol) && "skip".equals(param.handleInvalid))
  }


  test("sie transformation") {
    val data = TestingDataGenerator.spark.createDataFrame(
      Seq((0, "a"), (1, "b"), (2, "c"), (3, "a"), (4, "a"), (5, "c"))
    ).toDF("id", "category")
    val data2 = TestingDataGenerator.spark.createDataFrame(
      Seq((0, "a"), (1, "b"), (2, "e"))
    ).toDF("id", "category")
    val res = sie.train(param, data).transform(data2)

    // 输出结果 行数为2个
    res.show()
    assert(res.select(param.outputCol).rdd.count() == 2)
  }

}
